import numpy as np
import struct
from collections import defaultdict
import matplotlib.pyplot as plt
'''
first correct rate:0.8332
now correct rate:0.8413
'''
fig,ax=plt.subplots(
    nrows=2,
    ncols=5,
    sharex='all',
    sharey='all',
)


def normalize(data):  # 将图片像素二值化
    m, n = data.shape
    for i in range(m):
        for j in range(n):
            if data[i, j] != 0:  # 位置有像素即为1
                data[i, j] = 1
            else:
                data[i, j] = 0
    return data

def read_data():
    path = [r'Minist/t10k-images.idx3-ubyte',
            r'Minist/t10k-labels.idx1-ubyte',
            r'Minist/train-images.idx3-ubyte',
            r'Minist/train-labels.idx1-ubyte']  # 4个数据集路径
    data = defaultdict(dict)  # 声明一个空的字典 存放映射结果
    for i in range(0, 4):
        file = open(path[i], 'rb')  # 打开数据集
        f = file.read()
        file.close()
        # 以下进行数据的转化，图片、标签转化为数字信息
        if i % 2 == 0:  # 图片
            img_index = struct.calcsize('>IIII')
            _, size, row, column = struct.unpack('>IIII', f[:img_index])
            imgs = struct.unpack_from(str(size * row * column) + 'B', f, img_index)
            #一张图片数据存储为一行
            imgs = np.reshape(imgs, (size, row * column)).astype(np.float32)
            imgs = normalize(imgs)
            if i == 0:
                key = 'test'
            else:
                key = 'train'
            data[key]['images'] = imgs
        else:  # 标签
            label_index = struct.calcsize('>II')
            _, size = struct.unpack('>II', f[:label_index])
            labels = struct.unpack_from(str(size) + 'B', f, label_index)
            labels = np.reshape(labels, (size,))
            tmp = np.zeros((size, np.max(labels) + 1))
            tmp[np.arange(size), labels] = 1
            labels = tmp
            if i == 1:
                key = 'test'
            else:
                key = 'train'
            data[key]['labels'] = labels
    return data


def train(data):  # 训练模型
    '''

    :param data:
    :return:sum 10
    shape 10   28*28
    '''
    imgs = data['train']['images']  # 打开训练集
    labels = data['train']['labels']
    # n为训练集的训练个数（6W） dimsnum为转化后保存图片的所需的0 1位数 labelnum为标签的个数（10个 0-9）
    n, dimsnum = imgs.shape
    n, labelnum = labels.shape
    # 初始化sum和shape
    sum = np.zeros(labelnum)#sum记录着数字i在训练集中出现的个数
    shape = np.zeros((labelnum, dimsnum))
    for i in range(2):
        for j in range(5):
            global  ax
            if i==0:
                ax[i][j].imshow(imgs[j].reshape(28,28))
            else:
                ax[i][j].imshow(imgs[j+5].reshape(28,28))
    for i in range(0, n):  # n个数据
        # 找出第i个中最大的下标 labels[i][j]=1代表第i张图片为数字j 其余为0
        pos = np.argmax(labels[i])
        # 该数字出现次数+1
        sum[pos] = sum[pos] + 1
        for j in range(0, dimsnum):  # 图片转化后的n个维度
            shape[pos][j] = shape[pos][j] + imgs[i][j]  # shape[pos][j]代表数字pos第j维1的和
    for i in range(0, labelnum):
        for j in range(0, dimsnum):
            #shape[i][j] = (shape[i][j] + 1) / (sum[i] + dimsnum)  # 将shape转化为概率
            shape[i][j] = (shape[i][j] + 1) / (sum[i] + 2)
    #sum = (sum+1) / (n+labelnum)  # 将sum转化为概率
    sum = (sum + 0) / (n + 0)
    np.save('sum',sum)
    np.save('shape',shape)
    return sum, shape


def test(data, sum, shape):  # 测试模型
    imgs = data['test']['images']  # 打开测试集
    labels = data['test']['labels']
    # n为测试集个数 dimsnum为测试图片转化后的位数（维度） labelnum标签个数（0-9十个）
    n, dimsnum = imgs.shape
    n, labelnum = labels.shape
    correct = 0  # 测试时根据训练模型准确的个数
    for i in range(0, n):
        pos = np.argmax(labels[i])  # 图片对应的数字pos
        maxx = 0  # 0-9中的最大概率
        vpos = -1  # 0-9中的对应最大概率的数字
        for k in range(0, labelnum):  # 数字0-9
            ans = 1
            for j in range(0, dimsnum):  # dimsnum维
                if imgs[i][j] == 1:  # 第i个图片的第j维对应为1 '出现此概率'
                    ans *= shape[k][j]
                else:  # 未出现
                    ans *= (1 - shape[k][j])
            ans *= sum[k]  # 乘数字k在训练集中的概率 理想状态下为0.1
            if ans > maxx:
                maxx = ans
                vpos = k
        if vpos == pos:  # 当找到的数字和测试集的数字相同 测试成功数+1
            correct += 1
    return correct, n  # 返回成功数和测试集总数


if __name__ == '__main__':
    data = read_data()
    sum, shape = train(data)
    #correct, num = test(data, sum, shape)
    #print("测试正确率为：", correct / num)
    plt.show()
